#!env python
#-*- coding: utf-8 -*-


def generate(prompt,
             uri='http://localhost:11434/api/generate',
             model='deepseek-r1:latest',
             stream=True):
    import requests
    import json
    headers = { "Content-Type": "application/x-www-form-urlencoded" }
    data = json.dumps({ "model": model, "prompt": prompt, "stream": stream })
    response = requests.post(uri, headers=headers, data=data)
    if stream:
        # 按块迭代响应内容
        for chunk in response.iter_content(chunk_size=1024):
            text = chunk.decode("utf-8")
            print(chunk)
    return response


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('prompt', help='user input prompt')
    parser.add_argument('-u', '--uri', help='endpoint uri',
                        default='http://localhost:11434/api/generate')
    parser.add_argument('-m', '--model', help='model name',
                        default='deepseek-r1:latest')
    args = parser.parse_args()
    response = generate(args.prompt, args.uri, args.model)
    print(response)
